# Nuclear Engine Maya Mesh Export
# Code by Rockie X.Lee

# Usage:
#   import maya.cmds as cmds
#   cmds.loadPlugin('NEPATH/NEMayaMeshExport.py')
#   cmds.neMayaMeshExport()

import sys
import maya.OpenMaya as OpenMaya
import maya.OpenMayaMPx as OpenMayaMPx


class SubMesh():
    def __init__(self):
        self.__space = OpenMaya.MSpace.kWorld
        self.__name = "<>"
        self.__materialName = "<>"

        self.__vertexArray = OpenMaya.MFloatPointArray()
        self.__normalArray = OpenMaya.MFloatVectorArray()
        self.__tangentArray = OpenMaya.MFloatVectorArray()
        self.__colorArray = OpenMaya.MColorArray()
        self.__texCoordUArray = OpenMaya.MFloatArray()
        self.__texCoordVArray = OpenMaya.MFloatArray()

        self.__vertexTempArray = OpenMaya.MFloatPointArray()
        self.__normalTempArray = OpenMaya.MFloatVectorArray()
        self.__tangentTempArray = OpenMaya.MFloatVectorArray()
        self.__colorTempArray = OpenMaya.MColorArray()
        self.__texCoordUTempArray = OpenMaya.MFloatArray()
        self.__texCoordVTempArray = OpenMaya.MFloatArray()

        self.__hasVertex = False
        self.__hasNormal = False
        self.__hasTangent = False
        self.__hasColor = False
        self.__hasTexCoord = False

        self.__vertexIndexArray = []
        self.__normalIndexArray = []
        self.__tangentIndexArray = []
        self.__colorIndexArray = []
        self.__texCoordIndexArray = []

    def setData(self, dagPath):
        fnMesh = OpenMaya.MFnMesh(dagPath)
        self.__name = fnMesh.name()

        if fnMesh.numVertices() > 0:
            fnMesh.getPoints(self.__vertexTempArray, self.__space)
            self.__hasVertex = True

        if fnMesh.numNormals() > 0:
            fnMesh.getNormals(self.__normalTempArray, self.__space)
            self.__hasNormal = True

        if fnMesh.numUVs() > 0:
            fnMesh.getTangents(self.__tangentTempArray, self.__space)
            self.__hasTangent = True

        if fnMesh.numColors() > 0:
            fnMesh.getColors(self.__colorTempArray)
            self.__hasColor = True

        if fnMesh.numUVs() > 0:
            fnMesh.getUVs(self.__texCoordUTempArray, self.__texCoordVTempArray)
            self.__hasTexCoord = True

        tangentID = 0;
        iptr = OpenMaya.MScriptUtil().asIntPtr()
        for i in range(fnMesh.numPolygons()):
            numVertex = fnMesh.polygonVertexCount(i)
            vertexIndex = OpenMaya.MIntArray()
            fnMesh.getPolygonVertices(i, vertexIndex)
            self.__vertexIndexArray.append(vertexIndex)

            if self.__hasNormal:
                normalIndex = OpenMaya.MIntArray()
                fnMesh.getFaceNormalIds(i, normalIndex)
                self.__normalIndexArray.append(normalIndex)

            if self.__hasTangent:
                tangentIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    tangentIndex.append(tangentID)
                    tangentID += 1
                self.__tangentIndexArray.append(tangentIndex)

            if self.__hasColor:
                colorIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    fnMesh.getFaceVertexColorIndex(i, j, iptr)
                    colorIndex.append(OpenMaya.MScriptUtil.getInt(iptr))
                self.__colorIndexArray.append(colorIndex)

            if self.__hasTexCoord:
                texCoordIndex = OpenMaya.MIntArray()
                for j in range(numVertex):
                    fnMesh.getPolygonUVid(i, j, iptr)
                    texCoordIndex.append(OpenMaya.MScriptUtil.getInt(iptr))
                self.__texCoordIndexArray.append(texCoordIndex)

    def convertData(self):
        indexMap = dict()
        numIndex = 0
        for i in range(len(self.__vertexIndexArray)):
            for j in range(3):
                vertex = self.__vertexTempArray[self.__vertexIndexArray[i][j]]
                index = "V" + str(vertex.x) + "&" + str(vertex.y) + "&" + str(vertex.z)

                if self.__hasNormal:
                    normal = self.__normalTempArray[self.__normalIndexArray[i][j]]
                    index += "N" + str(normal.x) + "&" + str(normal.y) + "&" + str(normal.z)

                if self.__hasTangent:
                    tangent = self.__tangentTempArray[self.__tangentIndexArray[i][j]]
                    index += "T" + str(tangent.x) + "&" + str(tangent.y) + "&" + str(tangent.z)

                if self.__hasColor:
                    color = self.__colorTempArray[self.__colorIndexArray[i][j]]
                    index += "C" + str(color.r) + "&" + str(color.g) + "&" + str(color.b) + "&" + str(color.a)

                if self.__hasTexCoord:
                    texCoordU = self.__texCoordUTempArray[self.__texCoordIndexArray[i][j]]
                    texCoordV = self.__texCoordVTempArray[self.__texCoordIndexArray[i][j]]
                    index += "UV" + str(texCoordU) + "&" + str(texCoordV)

                if indexMap.has_key(index):
                    self.__vertexIndexArray[i][j] = indexMap[index]
                    """
                    if self.__hasNormal:
                        self.__normalIndexArray[i][j] = indexMap[index]
                    if self.__hasTangent:
                        self.__tangentIndexArray[i][j] = indexMap[index]
                    if self.__hasColor:
                        self.__colorIndexArray[i][j] = indexMap[index]
                    if self.__hasTexCoord:
                        self.__texCoordIndexArray[i][j] = indexMap[index]
                    """
                else:
                    indexMap[index] = numIndex
                    self.__vertexArray.append(vertex)
                    self.__vertexIndexArray[i][j] = numIndex
                    if self.__hasNormal:
                        self.__normalArray.append(normal)
                        #self.__normalIndexArray[i][j] = numIndex
                    if self.__hasTangent:
                        self.__tangentArray.append(tangent)
                        #self.__tangentIndexArray[i][j] = numIndex
                    if self.__hasColor:
                        self.__colorArray.append(color)
                        #self.__colorIndexArray[i][j] = numIndex
                    if self.__hasTexCoord:
                        self.__texCoordUArray.append(texCoordU)
                        self.__texCoordVArray.append(texCoordV)
                        #self.__texCoordIndexArray[i][j] = numIndex
                    numIndex += 1

    def writeData(self, dir):
        file = open(dir + self.__name + ".submeshpy", "w");
        file.write("Name = " + self.__name + "\n")
        file.write("MaterialName = " + self.__materialName + "\n")
        file.write("NumVertex = " + str(self.__vertexArray.length()) + "\n")
        file.write("NumNormal = " + str(self.__normalArray.length()) + "\n")
        file.write("NumTangent = " + str(self.__tangentArray.length()) + "\n")
        file.write("ColorSize = " + str(16) + "\n")
        file.write("NumColor = " + str(self.__colorArray.length()) + "\n")
        file.write("NumSecondaryColor = " + str(0) + "\n")
        file.write("TexCoordSize = " + str(8) + "\n")
        file.write("NumTexCoord = " + str(self.__texCoordUArray.length()) + "\n")
        file.write("NumBlendCount = " + str(4) + "\n")
        file.write("NumBlendWeight = " + str(len(self.__jointsWeightArray)) + "\n")
        file.write("NumBlenIndex = " + str(len(self.__jointsIndexArray)) + "\n")
        file.write("IndexSize = " + str(4) + "\n")
        file.write("NumIndex = " + str(len(self.__vertexIndexArray) * 3) + "\n")
        file.write("NumFace = " + str(len(self.__vertexIndexArray)) + "\n")

        for i in range(self.__vertexArray.length()):
            vertex = self.__vertexArray[i]
            file.write("V = " + str(vertex.x) + " " + str(vertex.y) + " " + str(vertex.z) + "\n")

        for i in range(self.__normalArray.length()):
            normal = self.__normalArray[i]
            file.write("N = " + str(normal.x) + " " + str(normal.y) + " " + str(normal.z) + "\n")

        for i in range(self.__tangentArray.length()):
            tangent = self.__tangentArray[i]
            file.write("T = " + str(tangent.x) + " " + str(tangent.y) + " " + str(tangent.z) + "\n")

        for i in range(self.__colorArray.length()):
            color = self.__colorArray[i]
            file.write("C = " + str(color.r) + " " + str(color.g) + " " + str(color.b) + " " + str(color.a) + "\n")

        for i in range(self.__texCoordUArray.length()):
            file.write("UV = " + str(self.__texCoordUArray[i]) + " " + str(self.__texCoordVArray[i]) + "\n")

        for i in range(len(self.__vertexIndexArray)):
            file.write("F = " + str(self.__vertexIndexArray[i][0]) + " " + str(self.__vertexIndexArray[i][1]) + " " + str(self.__vertexIndexArray[i][2]) + "\n")


neMayaMeshPluginCmdName = "neMayaMeshExport"

# Command
class NEMayaMeshExportCommand(OpenMayaMPx.MPxCommand):
    def __init__(self):
        OpenMayaMPx.MPxCommand.__init__(self)
        self.__dir = ""
        self.__name = "Mesh"
        self.__textureName = "<>"
        self.__bumpMapName = "<>"
        self.__subMeshArray = []

    def doIt(self, argList):
        print "[Begin] Nuclear Engine Maya Mesh Export===================================="
        self.exportMesh()
        print "[Success] Nuclear Engine Maya Mesh Export=================================="

    def exportMesh(self):
        # Create a selection list and iterator
        selectList = OpenMaya.MSelectionList()
        OpenMaya.MGlobal.getActiveSelectionList(selectList)
        # print selectList.length()
        iter = OpenMaya.MItSelectionList(selectList)
        dagIter = OpenMaya.MItDag(OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)

        # Save the dependency nodes in a list
        dependNode = OpenMaya.MObject()
        while not iter.isDone():
            iter.getDependNode(dependNode)
            objectPath = OpenMaya.MDagPath()
            iter.getDagPath(objectPath)
            dagIter.reset(objectPath.node(), OpenMaya.MItDag.kDepthFirst, OpenMaya.MFn.kInvalid)
            while not dagIter.isDone():
                dagPath = OpenMaya.MDagPath()
                dagIter.getPath(dagPath)
                if dagPath.hasFn(OpenMaya.MFn.kMesh) and dagPath.hasFn(OpenMaya.MFn.kTransform):
                    pass
                elif dagPath.hasFn(OpenMaya.MFn.kMesh):
                    fnMesh = OpenMaya.MFnMesh(dagPath)
                    self.__subMeshArray.append(fnMesh.name())
                    subMesh = SubMesh()
                    subMesh.setData(dagPath)
                    subMesh.convertData()
                    subMesh.writeData(self.__dir)
                dagIter.next()
            iter.next()
        self.writeData(self.__dir)

    def writeData(self, dir):
        file = open(dir + self.__name + ".meshpy", "w");
        file.write("Name = " + self.__name + "\n")
        file.write("TextureName = " + self.__textureName + "\n")
        file.write("BumpMapName = " + self.__bumpMapName + "\n")
        file.write("NumSubMeshs = " + str(len(self.__subMeshArray)) + "\n")
        for i in range(len(self.__subMeshArray)):
            file.write("SM = " + self.__subMeshArray[i] + "\n")
        file.close()

# Creator
def cmdCreator():
    return OpenMayaMPx.asMPxPtr(NEMayaMeshExportCommand())


# Syntax creator
def syntaxCreator():
    syntax = OpenMaya.MSyntax()
    return syntax


# Initialize the plug-in
def initializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.registerCommand(neMayaMeshPluginCmdName, cmdCreator, syntaxCreator)
    except:
        sys.stderr.write("Failed to register command: %s\n" % neMayaMeshPluginCmdName)
        raise


# Uninitialize the plug-in
def uninitializePlugin(mayaObject):
    mayaPlugin = OpenMayaMPx.MFnPlugin(mayaObject)
    try:
        mayaPlugin.deregisterCommand(neMayaMeshPluginCmdName)
    except:
        sys.stderr.write("Failed to unregister command: %s\n" % neMayaMeshPluginCmdName)
        raise

